package org.finesys.chat.core.base.embedding.store.redis;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import org.finesys.common.core.util.ValidationUtil;

import redis.clients.jedis.json.Path2;
import redis.clients.jedis.search.schemafields.SchemaField;
import redis.clients.jedis.search.schemafields.TextField;
import redis.clients.jedis.search.schemafields.VectorField;

public class RedisSchema {

    public static final String SCORE_FIELD_NAME = "vector_score";
    public static final String JSON_KEY = "$";
    public static final Path2 JSON_SET_PATH = Path2.of("$");
    public static final String JSON_PATH_PREFIX = "$.";
    private static final VectorField.VectorAlgorithm DEFAULT_VECTOR_ALGORITHM;
    private static final MetricType DEFAULT_METRIC_TYPE;
    private final String indexName;
    private final String prefix;
    private final String vectorFieldName;
    private final String scalarFieldName;
    private final Map<String, SchemaField> metadataConfig;
    private final VectorField.VectorAlgorithm vectorAlgorithm;
    private final Integer dimension;
    private final MetricType metricType;

    static {
        DEFAULT_VECTOR_ALGORITHM = VectorField.VectorAlgorithm.FLAT;
        DEFAULT_METRIC_TYPE = MetricType.COSINE_DISTANCE;
    }

    RedisSchema(String indexName, String prefix, String vectorFieldName, String scalarFieldName, VectorField.VectorAlgorithm vectorAlgorithm, Integer dimension, MetricType metricType, Map<String, SchemaField> metadataConfig) {
        ValidationUtil.ensureTrue(prefix.endsWith(":"), "Prefix should end with a ':'");
        this.indexName = indexName;
        this.prefix = prefix;
        this.vectorFieldName = vectorFieldName;
        this.scalarFieldName = scalarFieldName;
        this.vectorAlgorithm = vectorAlgorithm;
        this.dimension = dimension;
        this.metricType = metricType;
        this.metadataConfig = metadataConfig;
    }

    SchemaField[] toSchemaFields() {
        Map<String, Object> vectorMap = new HashMap<>();
        vectorMap.put(SchemaConstants.DIMENSION, RedisSchema.this.dimension);
        vectorMap.put(SchemaConstants.DISTANCE_METRIC, this.metricType.getValue());
        vectorMap.put(SchemaConstants.TYPE, "FLOAT32");
        vectorMap.put(SchemaConstants.INITIAL_CAP, 5);
        List<SchemaField> schemaFields = new ArrayList<>();
        schemaFields.add(TextField.of("$." + this.scalarFieldName).as(this.scalarFieldName).weight(1.0));
        schemaFields.add(VectorField.builder().fieldName("$." + this.vectorFieldName).algorithm(this.vectorAlgorithm).attributes(vectorMap).as(this.vectorFieldName).build());
        schemaFields.addAll(this.metadataConfig.values());
        return schemaFields.toArray(new SchemaField[0]);
    }

    String indexName() {
        return this.indexName;
    }

    String prefix() {
        return this.prefix;
    }

    String vectorFieldName() {
        return this.vectorFieldName;
    }

    String scalarFieldName() {
        return this.scalarFieldName;
    }

    VectorField.VectorAlgorithm vectorAlgorithm() {
        return this.vectorAlgorithm;
    }

    Integer dimension() {
        return this.dimension;
    }

    MetricType metricType() {
        return this.metricType;
    }

    Map<String, SchemaField> metadataConfig() {
        return this.metadataConfig;
    }


    public static Builder builder() {
        return new Builder();
    }

    public static class Builder {
        private String indexName;
        private String prefix = "embedding:";
        private String vectorFieldName = "vector";
        private String scalarFieldName = "text";
        private Map<String, SchemaField> metadataConfig = new HashMap<>();
        private VectorField.VectorAlgorithm vectorAlgorithm;
        private Integer dimension;
        private MetricType metricType;

        public Builder() {
            this.metricType = RedisSchema.DEFAULT_METRIC_TYPE;
            this.vectorAlgorithm = RedisSchema.DEFAULT_VECTOR_ALGORITHM;
        }

        public Builder indexName(String indexName) {
            this.indexName = indexName;
            return this;
        }

        public Builder prefix(String prefix) {
            this.prefix = prefix;
            return this;
        }

        public Builder vectorFieldName(String vectorFieldName) {
            this.vectorFieldName = vectorFieldName;
            return this;
        }

        public Builder scalarFieldName(String scalarFieldName) {
            this.scalarFieldName = scalarFieldName;
            return this;
        }

        public Builder metadataConfig(Map<String, SchemaField> metadataConfig) {
            this.metadataConfig = metadataConfig;
            return this;
        }

        public Builder vectorAlgorithm(VectorField.VectorAlgorithm vectorAlgorithm) {
            this.vectorAlgorithm = vectorAlgorithm;
            return this;
        }

        public Builder dimension(Integer dimension) {
            this.dimension = dimension;
            return this;
        }

        public Builder metricType(MetricType metricType) {
            this.metricType = metricType;
            return this;
        }

        public RedisSchema build() {
            return new RedisSchema(this.indexName, this.prefix, this.vectorFieldName, this.scalarFieldName, this.vectorAlgorithm, this.dimension, this.metricType, this.metadataConfig);
        }
    }
}
